import numpy as np
import tmm
from concurrent.futures import ProcessPoolExecutor

# 定义材料的折射率作为常量
N_SIO2 = 1.45
N_SI3N4 = 2


def validate_parameters(num_layers, max_t, t_sample_num, freq_start, freq_end, num_freq):
    """验证输入参数的有效性"""
    if not isinstance(num_layers, int) or num_layers <= 0:
        raise ValueError("num_layers必须是正整数")
    if not isinstance(max_t, (int, float)) or max_t <= 0:
        raise ValueError("max_t必须是正数")
    if not isinstance(t_sample_num, int) or t_sample_num <= 0:
        raise ValueError("t_sample_num必须是正整数")
    if not isinstance(freq_start, (int, float)) or not isinstance(freq_end, (int, float)) or freq_start >= freq_end:
        raise ValueError("freq_start必须小于freq_end")
    if not isinstance(num_freq, int) or num_freq <= 0:
        raise ValueError("num_freq必须是正整数")


def generate_t_sample_array_elem(i, t_sample, t_sample_num, num_layers):
    binary_repr = np.base_repr(i, t_sample_num).zfill(num_layers)
    t_sample_array_elem = np.array([t_sample[int(x)] for x in binary_repr])
    return t_sample_array_elem


def calculate_transmission_parallel(t_sample_array_elem, n_array, unit, n_top, n_bottom, freq_start, freq_end,
                                    num_freq):
    """计算传输特性，用于并行处理"""
    try:
        _, tran, _ = np.abs(tmm.TMM(t_sample_array_elem, n_array, unit=unit, n_top=n_top, n_bottom=n_bottom)
                            .calculate_spectrum(freq_start, freq_end, num_freq)) ** 2
        return tran
    except Exception as e:
        print(f"计算传输特性时发生异常: {e}")
        return None


def data_generator(num_layers, max_t, t_sample_num, freq_start, freq_end, num_freq, unit='nm/THz', n_top=1, n_bottom=1):
    validate_parameters(num_layers, max_t, t_sample_num, freq_start, freq_end, num_freq)

    n_sio2 = np.ones(num_layers // 2) * N_SIO2
    n_si3n4 = np.ones(num_layers // 2) * N_SI3N4
    n_array = np.column_stack((n_sio2, n_si3n4)).flatten()
    if num_layers % 2 != 0:
        n_array = np.append(n_array, N_SIO2)

    t_sample = np.linspace(10, max_t, t_sample_num)
    data_num = t_sample_num ** num_layers

    with ProcessPoolExecutor() as executor1:
        t_sample_array = list(executor1.map(generate_t_sample_array_elem,
                                            np.arange(data_num, dtype=int),
                                            np.repeat(t_sample[np.newaxis, :], data_num, axis=0),
                                            np.ones(data_num, dtype=int) * t_sample_num,
                                            np.ones(data_num, dtype=int) * num_layers))

    with ProcessPoolExecutor() as executor2:
        tran_array = list(executor2.map(calculate_transmission_parallel,
                                        t_sample_array,
                                        np.repeat(n_array[np.newaxis, :], data_num, axis=0),
                                        [unit] * data_num,
                                        np.ones(data_num) * n_top,
                                        np.ones(data_num) * n_bottom,
                                        np.ones(data_num) * freq_start,
                                        np.ones(data_num) * freq_end,
                                        np.ones(data_num, dtype=int) * num_freq))

    # 用try-except捕获np.save可能的异常
    try:
        np.save('train_set/t_array.npy', np.array(t_sample_array))
        np.save('train_set/tran_array.npy', np.array(tran_array))
    except Exception as e:
        print(f"保存数组时发生异常: {e}")

    return


if __name__ == '__main__':
    data_generator(10, 100, 3, 450, 750, 200)
